package io.github.wanggit.antrpc.client.zk.lb;

import io.github.wanggit.antrpc.client.zk.zknode.NodeHostEntity;
import io.github.wanggit.antrpc.commons.bean.RegisterBeanHelper;
import io.github.wanggit.antrpc.commons.bean.RegisterBeanMethod;
import io.github.wanggit.antrpc.commons.bean.SubscribeNode;
import io.github.wanggit.antrpc.commons.lb.LoadBalancerDTO;
import io.github.wanggit.antrpc.commons.rate.IResilience4jRateLimiter;
import io.github.wanggit.antrpc.commons.utils.Utils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.ClassUtils;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/** 负载均衡容器，每一个方法对应一个负载均衡器 */
@Slf4j
public class LoadBalancerContainer implements ILoadBalancerContainer {

    private final ConcurrentHashMap<String, ILoadBalancer> loadBalancers =
            new ConcurrentHashMap<>();

    private final LoadBalancerHelper loadBalancerHelper;
    private final IResilience4jRateLimiter clientResilience4jRateLimiter;

    public LoadBalancerContainer(
            LoadBalancerHelper loadBalancerHelper,
            IResilience4jRateLimiter clientResilience4jRateLimiter) {
        this.loadBalancerHelper = loadBalancerHelper;
        this.clientResilience4jRateLimiter = clientResilience4jRateLimiter;
    }

    @Override
    public HashMap<String, ILoadBalancer> loadBalancersSnapshot() {
        return new HashMap<>(loadBalancers);
    }

    @Override
    public void addLoadBalancer(LoadBalancerDTO loadBalancerDTO) {
        internalSaveOrUpdateLoadBalancer(loadBalancerDTO);
    }

    @Override
    public void updateLoadBalancer(LoadBalancerDTO loadBalancerDTO) {
        internalSaveOrUpdateLoadBalancer(loadBalancerDTO);
    }

    @Override
    public void deleteLoadBalancer(String interfaceName) {
        Set<String> keys = new HashSet<>();
        for (Map.Entry<String, ILoadBalancer> entry : loadBalancers.entrySet()) {
            if (entry.getKey().startsWith(interfaceName + "#")) {
                keys.add(entry.getKey());
            }
        }
        for (String key : keys) {
            loadBalancers.remove(key);
        }
    }

    @Override
    public ILoadBalancer getLoadBalancer(String classMethodFullName) {
        ILoadBalancer loadBalancer = loadBalancers.get(classMethodFullName);
        if (null == loadBalancer) {
            synchronized (classMethodFullName.intern()) {
                if (!loadBalancers.containsKey(classMethodFullName)) {
                    loadBalancers.put(classMethodFullName, loadBalancerHelper.getLoadBalancer());
                }
            }
        }
        return loadBalancers.get(classMethodFullName);
    }

    @Override
    public void addHost(String classMethodFullName, NodeHostEntity host) {
        getLoadBalancer(classMethodFullName).addHost(host);
    }

    @Override
    public void updateHost(String classMethodFullName, NodeHostEntity host) {
        getLoadBalancer(classMethodFullName).updateHost(host);
    }

    @Override
    public void deleteHost(String classMethodFullName, NodeHostEntity host) {
        getLoadBalancer(classMethodFullName).deleteHost(host);
    }

    @Override
    public Map<String, List<NodeHostEntity>> getAllNodeHostEntities() {
        Map<String, List<NodeHostEntity>> datas = new HashMap<>();
        loadBalancers.forEach((key, value) -> datas.put(key, value.getAllNodeHostEntities()));
        return datas;
    }

    @Override
    public void updateHostRateLimit(SubscribeNode subscribeNode) {
        Map<String, RegisterBeanMethod> methodMap = subscribeNode.getMethodMap();
        if (null != methodMap) {
            methodMap.forEach(
                    (key, value) -> {
                        // RegisterBeanMethod中的className可能是实现类，把它替换为接口
                        value.setClassName(subscribeNode.getClassName());
                        // 检查并重新生成频控管理器
                        clientResilience4jRateLimiter.checkOrChangeEventCountCircuitBreaker(value);
                    });
        }
    }

    private void internalSaveOrUpdateLoadBalancer(LoadBalancerDTO loadBalancerDTO) {
        Class<?> clazz = classForName(loadBalancerDTO.getClazz());
        if (null != clazz) {
            RegisterBeanMethod[] registerBeanMethods =
                    RegisterBeanHelper.findMethodsWithoutJavaBasicMethods(clazz);
            for (RegisterBeanMethod registerBeanMethod : registerBeanMethods) {
                ILoadBalancer loadBalancer = newLoadBalancerInstance(loadBalancerDTO.getLb());
                if (null != loadBalancer) {
                    String classMethodFullName =
                            Utils.classMethodFullName(
                                    loadBalancerDTO.getClazz(), registerBeanMethod.toString());
                    ILoadBalancer oldLoadBalancer = getLoadBalancer(classMethodFullName);
                    ((AbstractLoadBalancer) oldLoadBalancer)
                            .copyEntitiesTo((AbstractLoadBalancer) loadBalancer);
                    loadBalancer.setWeights(loadBalancerDTO.getWeights());
                    loadBalancers.put(classMethodFullName, loadBalancer);
                }
            }
        }
    }

    private Class<?> classForName(String className) {
        try {
            return ClassUtils.forName(className, LoadBalancerContainer.class.getClassLoader());
        } catch (ClassNotFoundException e) {
            if (log.isWarnEnabled()) {
                log.warn(className + " class loading failed!");
            }
        }
        return null;
    }

    private ILoadBalancer newLoadBalancerInstance(String loadBalancerClazz) {
        if (null == loadBalancerClazz) {
            return loadBalancerHelper.getLoadBalancer();
        }
        try {
            Class<?> aClass =
                    ClassUtils.forName(
                            loadBalancerClazz, LoadBalancerContainer.class.getClassLoader());
            Object instance = aClass.newInstance();
            return (ILoadBalancer) instance;
        } catch (Exception e) {
            if (log.isErrorEnabled()) {
                log.error(loadBalancerClazz + " instance error. ", e);
            }
            return loadBalancerHelper.getLoadBalancer();
        }
    }
}
